#!/usr/bin/env python3
"""
文本处理模型管理器模块

统一管理所有文本处理相关的模型，包括StructBERT和PALM2.0
使用与funasr_wss_server相同的模型加载机制
"""

import logging
from typing import Any, Optional

# 检查ModelScope可用性（文本处理模型需要）
try:
    from modelscope.pipelines import pipeline
    MODELSCOPE_AVAILABLE = True
except ImportError:
    logging.warning("ModelScope未安装，文本处理功能不可用")
    MODELSCOPE_AVAILABLE = False

from ...config.funasr_config import ModelConfig
from .palm2_manager import Palm2Manager
from .structbert_manager import StructBertManager

logger = logging.getLogger(__name__)

class UnifiedTextProcessingManager:
    """统一文本处理模型管理器"""

    def __init__(self):
        self.structbert_manager = None
        self.palm2_manager = None
        self.models_loaded = False

    async def initialize(self) -> bool:
        """初始化所有文本处理模型"""
        try:
            if not MODELSCOPE_AVAILABLE:
                logger.error("ModelScope未安装，无法加载文本处理模型")
                logger.info("提示：请安装ModelScope：pip install modelscope")
                return False

            # 初始化StructBERT管理器
            self.structbert_manager = StructBertManager()

            # 初始化PALM2管理器
            self.palm2_manager = Palm2Manager()

            # 加载模型
            structbert_success = await self.structbert_manager.load_model()
            palm2_success = await self.palm2_manager.load_model()

            self.models_loaded = structbert_success or palm2_success

            if structbert_success:
                logger.info("StructBERT模型加载成功")
            else:
                logger.warning("StructBERT模型加载失败")

            if palm2_success:
                logger.info("PALM2.0模型加载成功")
            else:
                logger.warning("PALM2.0模型加载失败")

            return self.models_loaded

        except Exception as e:
            logger.error(f"文本处理模型初始化失败: {e}")
            return False

    def get_structbert_manager(self) -> Optional[StructBertManager]:
        """获取StructBERT管理器"""
        return self.structbert_manager

    def get_palm2_manager(self) -> Optional[Palm2Manager]:
        """获取PALM2管理器"""
        return self.palm2_manager

    def is_ready(self) -> bool:
        """检查模型是否就绪"""
        return self.models_loaded and (
            (self.structbert_manager and self.structbert_manager.model_loaded) or
            (self.palm2_manager and self.palm2_manager.model_loaded)
        )

# 全局实例
_text_processing_manager = None

async def get_text_processing_manager() -> UnifiedTextProcessingManager:
    """获取全局文本处理管理器实例"""
    global _text_processing_manager
    if _text_processing_manager is None:
        _text_processing_manager = UnifiedTextProcessingManager()
        await _text_processing_manager.initialize()
    return _text_processing_manager

__all__ = ["UnifiedTextProcessingManager", "get_text_processing_manager", "StructBertManager", "Palm2Manager"]
